#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
HanLP RESTful API服务器
提供HanLP的Web API接口
"""

import os
import sys
import json
import logging
import argparse
import time
from functools import wraps
from flask import Flask, request, jsonify, render_template
from flask_cors import CORS

import hanlp

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S"
)
logger = logging.getLogger("hanlp-server")

# 创建Flask应用
app = Flask(__name__)
CORS(app)  # 启用跨域支持

# 全局变量存储模型
models = {}

# 计时装饰器
def timer(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        start = time.time()
        result = f(*args, **kwargs)
        end = time.time()
        logger.info(f"{f.__name__} 执行时间: {end - start:.4f}秒")
        return result
    return wrapper

@timer
def load_models(use_mtl=False):
    """加载HanLP模型"""
    try:
        logger.info("开始加载HanLP模型...")
        
        # 加载分词模型
        logger.info("正在加载分词模型...")
        models['tok'] = hanlp.load(hanlp.pretrained.tok.FINE_ELECTRA_SMALL_ZH)
        
        # 加载词性标注模型
        logger.info("正在加载词性标注模型...")
        models['pos'] = hanlp.load(hanlp.pretrained.pos.CTB9_POS_ELECTRA_SMALL)
        
        # 加载命名实体识别模型
        logger.info("正在加载命名实体识别模型...")
        models['ner'] = hanlp.load(hanlp.pretrained.ner.MSRA_NER_ELECTRA_SMALL_ZH)
        
        # 加载依存句法分析模型
        logger.info("正在加载依存句法分析模型...")
        models['dep'] = hanlp.load(hanlp.pretrained.dep.CTB7_BIAFFINE_DEP_ZH)
        
        # 加载多任务模型（可选）
        if use_mtl:
            logger.info("正在加载多任务模型...")
            models['mtl'] = hanlp.load(hanlp.pretrained.mtl.ALL_IN_ONE_ELECTRA_SMALL_ZH)
        
        logger.info("所有模型加载完成")
        return True
    except Exception as e:
        logger.error(f"模型加载失败: {str(e)}")
        return False

# 路由：HTML客户端界面
@app.route('/', methods=['GET'])
def index():
    """提供HTML客户端界面"""
    if os.path.exists(os.path.join(app.root_path, 'templates', 'index.html')):
        return render_template('index.html')
    else:
        return jsonify({
            'status': 'ok',
            'message': 'HanLP RESTful API服务已启动',
            'endpoints': {
                '/api/tokenize': '文本分词',
                '/api/pos': '词性标注',
                '/api/ner': '命名实体识别',
                '/api/dep': '依存句法分析',
                '/api/mtl': '多任务处理（如果启用）'
            }
        })

# 路由：分词
@app.route('/api/tokenize', methods=['POST'])
@timer
def tokenize():
    """分词API"""
    try:
        data = request.json
        if not data or 'text' not in data:
            return jsonify({'status': 'error', 'message': '缺少必要的文本字段'}), 400
        
        text = data['text']
        logger.info(f"分词请求: {text[:100]}...")
        
        tokens = models['tok'](text)
        
        return jsonify({
            'status': 'ok',
            'message': '分词成功',
            'result': tokens
        })
    except Exception as e:
        logger.error(f"分词错误: {str(e)}")
        return jsonify({'status': 'error', 'message': f'处理过程中出错: {str(e)}'}), 500

# 路由：词性标注
@app.route('/api/pos', methods=['POST'])
@timer
def pos_tagging():
    """词性标注API"""
    try:
        data = request.json
        if not data or 'text' not in data:
            return jsonify({'status': 'error', 'message': '缺少必要的文本字段'}), 400
        
        text = data['text']
        logger.info(f"词性标注请求: {text[:100]}...")
        
        tokens = models['tok'](text)
        pos_tags = models['pos'](tokens)
        
        result = list(zip(tokens, pos_tags))
        
        return jsonify({
            'status': 'ok',
            'message': '词性标注成功',
            'result': {
                'tokens': tokens,
                'pos_tags': pos_tags,
                'paired': result
            }
        })
    except Exception as e:
        logger.error(f"词性标注错误: {str(e)}")
        return jsonify({'status': 'error', 'message': f'处理过程中出错: {str(e)}'}), 500

# 路由：命名实体识别
@app.route('/api/ner', methods=['POST'])
@timer
def ner():
    """命名实体识别API"""
    try:
        data = request.json
        if not data or 'text' not in data:
            return jsonify({'status': 'error', 'message': '缺少必要的文本字段'}), 400
        
        text = data['text']
        logger.info(f"命名实体识别请求: {text[:100]}...")
        
        tokens = models['tok'](text)
        entities = models['ner'](tokens)
        
        # 处理实体结果
        result = []
        for entity in entities:
            if len(entity) == 3:  # 确保实体包含起始位置、结束位置和类型
                start, end, entity_type = entity
                entity_text = ''.join(tokens[start:end])
                result.append({
                    'text': entity_text,
                    'type': entity_type,
                    'start': start,
                    'end': end
                })
        
        return jsonify({
            'status': 'ok',
            'message': '命名实体识别成功',
            'result': {
                'tokens': tokens,
                'entities': result
            }
        })
    except Exception as e:
        logger.error(f"命名实体识别错误: {str(e)}")
        return jsonify({'status': 'error', 'message': f'处理过程中出错: {str(e)}'}), 500

# 路由：依存句法分析
@app.route('/api/dep', methods=['POST'])
@timer
def dependency_parsing():
    """依存句法分析API"""
    try:
        data = request.json
        if not data or 'text' not in data:
            return jsonify({'status': 'error', 'message': '缺少必要的文本字段'}), 400
        
        text = data['text']
        logger.info(f"依存句法分析请求: {text[:100]}...")
        
        tokens = models['tok'](text)
        tree = models['dep'](tokens)
        
        # 将依存句法树转换为易于前端显示的格式
        result = []
        for i, (token, head, rel) in enumerate(zip(tokens, tree.arc_ls, tree.rel_ls)):
            result.append({
                'id': i + 1,
                'text': token,
                'head': head,  # 头节点的索引（从1开始）
                'relation': rel  # 关系类型
            })
        
        return jsonify({
            'status': 'ok',
            'message': '依存句法分析成功',
            'result': result
        })
    except Exception as e:
        logger.error(f"依存句法分析错误: {str(e)}")
        return jsonify({'status': 'error', 'message': f'处理过程中出错: {str(e)}'}), 500

# 路由：多任务处理
@app.route('/api/mtl', methods=['POST'])
@timer
def multitask():
    """多任务处理API"""
    if 'mtl' not in models:
        return jsonify({'status': 'error', 'message': '多任务模型未加载'}), 400
    
    try:
        data = request.json
        if not data or 'text' not in data:
            return jsonify({'status': 'error', 'message': '缺少必要的文本字段'}), 400
        
        text = data['text']
        logger.info(f"多任务处理请求: {text[:100]}...")
        
        result = models['mtl'](text)
        
        # 将结果转换为JSON可序列化格式
        serializable_result = {}
        for key, value in result.items():
            if hasattr(value, 'tolist'):  # 处理numpy数组
                serializable_result[key] = value.tolist()
            elif isinstance(value, list) and value and hasattr(value[0], 'tolist'):  # 处理numpy数组列表
                serializable_result[key] = [item.tolist() if hasattr(item, 'tolist') else item for item in value]
            else:
                # 尝试转换其他复杂对象
                try:
                    json.dumps({key: value})
                    serializable_result[key] = value
                except (TypeError, OverflowError):
                    serializable_result[key] = str(value)
        
        return jsonify({
            'status': 'ok',
            'message': '多任务处理成功',
            'result': serializable_result
        })
    except Exception as e:
        logger.error(f"多任务处理错误: {str(e)}")
        return jsonify({'status': 'error', 'message': f'处理过程中出错: {str(e)}'}), 500

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='启动HanLP RESTful API服务')
    parser.add_argument('--host', type=str, default='0.0.0.0', help='服务器主机地址')
    parser.add_argument('--port', type=int, default=8000, help='服务器端口号')
    parser.add_argument('--debug', action='store_true', help='启用调试模式')
    parser.add_argument('--multitask', action='store_true', help='加载多任务模型')
    
    args = parser.parse_args()
    
    logger.info(f"正在启动HanLP服务器 (host={args.host}, port={args.port}, debug={args.debug}, multitask={args.multitask})")
    
    # 加载模型
    if not load_models(use_mtl=args.multitask):
        logger.error("模型加载失败，服务无法启动")
        sys.exit(1)
    
    # 启动Flask应用
    app.run(host=args.host, port=args.port, debug=args.debug)

if __name__ == '__main__':
    main() 